package cn.fengtang.chatglm.data.domain.service;

import cn.fengtang.chatglm.data.domain.model.aggregates.ChatProcessAggregate;
import cn.fengtang.chatglm.model.ChatCompletionRequest;
import cn.fengtang.chatglm.model.ChatCompletionResponse;
import cn.fengtang.chatglm.model.Role;
import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.apache.commons.lang3.StringUtils;
import org.jetbrains.annotations.NotNull;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;

import javax.annotation.Nullable;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

@Slf4j
@Service
public class ChatService extends AbstractChatService {

    @Override
    protected void doMessageResponse(ChatProcessAggregate chatProcess, ResponseBodyEmitter emitter) throws Exception {
        // 1. 请求消息
        List<ChatCompletionRequest.Prompt> messages = chatProcess.getMessages().stream().map(entity -> ChatCompletionRequest.Prompt.builder().role(entity.getRole()).content(entity.getContent()).build()).collect(Collectors.toList());

        // 2. 构建请求
        ChatCompletionRequest request = new ChatCompletionRequest();
        // 创建一个列表用于存储所有 Prompt 对象
        List<ChatCompletionRequest.Prompt> promptList = new ArrayList<>();
        for (ChatCompletionRequest.Prompt message : messages) {
            request.setTools(new ArrayList<ChatCompletionRequest.Tool>() {
                private static final long serialVersionUID = -7988151926241837899L;

                {
                    add(ChatCompletionRequest.Tool.builder().type(ChatCompletionRequest.Tool.Type.web_search).webSearch(ChatCompletionRequest.Tool.WebSearch.builder().enable(true).searchQuery(message.getContent()).build()).build());
                }
            });

            // 创建 Prompt 对象并添加到列表中
            ChatCompletionRequest.Prompt prompt = ChatCompletionRequest.Prompt.builder().role(Role.user.getCode()).content(message.getContent()).build();
            promptList.add(prompt);
        }
        // 使用 setPrompt 方法设置 Prompt 列表
        request.setPrompt(promptList);

        // 请求
        openAiSession.completions(request, new EventSourceListener() {

            @Override
            public void onEvent(@NotNull EventSource eventSource, @Nullable String id, @Nullable String type, @NotNull String data) {
                log.info("收到数据块[{}],{}", data, System.currentTimeMillis());
                ChatCompletionResponse chatCompletionResponse = JSON.parseObject(data, ChatCompletionResponse.class);
                List<ChatCompletionResponse.Choice> choices = chatCompletionResponse.getChoices();
                for (ChatCompletionResponse.Choice chatChoice : choices) {
                    ChatCompletionResponse.Delta delta = chatChoice.getDelta();
                    String finishReason = chatChoice.getFinishReason();
                    if (StringUtils.isNoneBlank(finishReason) && "stop".equals(finishReason)) {
                        emitter.complete();
                        break;
                    }

                    // 发送信息
                    try {
                        Thread.sleep(10);
                        emitter.send(delta.getContent());
                    } catch (Exception e) {
                        throw new RuntimeException(e);
                    }

                }
            }
        });

    }
}

